from .neo4j_handler import Neo4jHandler
from .qwen_generator import generate_analysis
from .markdown_generator import generate_markdown, save_markdown
from .fault_matcher import FaultMatcher
from .ansible_executor import execute_solution_code, extract_solution1_code
from .knowledge_loader import load_script_by_code_id
from .graphrag_retriever import GraphRAGRetriever

def main(user_fault_input):
    # 1. 匹配标准故障类型
    matcher = FaultMatcher()
    standard_fault = matcher.match_standard_fault(user_fault_input)
    if not standard_fault:
        print("无法在知识图谱中匹配到标准故障类型")
        return
    print(f"用户输入：{user_fault_input} → 匹配到标准故障：{standard_fault}")

    # 2. 获取标准故障的详细信息
    fault_details = matcher.get_fault_details(standard_fault)
    if not fault_details:
        print("无法获取标准故障的详细信息")
        return

    # 3. 从知识图谱获取更多故障关联信息
    neo4j = Neo4jHandler()
    fault_info = neo4j.query_fault_info(standard_fault)

    # 4. 显示知识图谱查询结果
    print("===== 知识图谱查询结果 =====")
    print(f"故障名称：{fault_info.get('fault_name', '未知')}")
    print(f"根本原因实体：{fault_info.get('root_cause', {}).get('entity', '未知')}")
    print(f"修复代码ID：{fault_info.get('fix_code_id', '未配置')}")

    # 5. 查询上游故障信息
    dependencies = fault_info.get('dependencies', [])
    upstream_faults = []
    seen_faults = set()
    for dep in dependencies:
        upstream_entity = dep["upstream_entity"]
        current_upstream_faults = neo4j.query_upstream_fault_info(upstream_entity)
        for fault in current_upstream_faults:
            if fault["fault_name"] not in seen_faults:
                seen_faults.add(fault["fault_name"])
                upstream_faults.append(fault)
    print(f"上游关联故障数量：{len(upstream_faults)}")

    #-------------------------- GraphRAG检索增强知识 --------------------------
    print("\n===== GraphRAG增强检索 =====")
    try:
        # 初始化GraphRAG检索器
        rag_retriever = GraphRAGRetriever(graphrag_root_path="/root/graphrag-basic/ragtest")
        # 检索当前故障的增强知识
        enhanced_knowledge = rag_retriever.retrieve_enhanced_knowledge(standard_fault)
        # 将增强知识补充到fault_info中，供大模型使用
        if enhanced_knowledge:
            fault_info["enhanced_knowledge"] = enhanced_knowledge
        else:
            fault_info["enhanced_knowledge"] = "【GraphRAG增强知识】未检索到相关补充信息"
    except Exception as e:
        # 检索失败不阻断主流程，仅记录警告
        print(f"⚠️ GraphRAG检索初始化失败，跳过增强步骤：{str(e)}")
        fault_info["enhanced_knowledge"] = "【GraphRAG增强知识】检索模块初始化失败，未获取补充信息"
    # --------------------------------------------------------------------------------

    # 6. 调用大模型生成解决方案
    print("\n===== 生成故障分析报告 =====")
    analysis = generate_analysis(standard_fault, fault_details, fault_info, upstream_faults)

    # 7. 执行脚本和生成报告
    execution_result = ""
    script_content = None
    is_predefined = False
    if fault_info.get("fix_code_id"):
        script_content = load_script_by_code_id(fault_info["fix_code_id"])
        if script_content:
            print(f"✅ 成功加载预定义脚本")
            is_predefined = True
        else:
            print(f"❌ 预定义脚本加载失败")
    if not script_content:
        script_content = extract_solution1_code(analysis)
        if not script_content:
            execution_result = "❌ 无法获取修复脚本，未执行修复任务"
            print(execution_result)
    if script_content:
        print("\n===== 开始执行修复脚本 =====")
        try:
            execution_result = execute_solution_code(
                fault_name=standard_fault,
                shell_code=script_content,
                inventory="localhost"
            )
            print(f"✅ 脚本执行完成")
        except Exception as e:
            error_msg = f"❌ 执行过程错误：{str(e)}"
            execution_result = error_msg
            print(error_msg)
    md_content = generate_markdown(
        standard_fault,
        analysis,
        execution_result,
        script_content
    )
    save_markdown(md_content)

if __name__ == "__main__":
    # 保留测试入口
    test_fault = "CPU 使用率过高（负载过高）"
    print(f"测试模式：模拟检测到故障 -> {test_fault}")
    main(test_fault)